package main

import (
	"encoding/csv"
	"fmt"
	"os"
	"path/filepath"

	"github.com/xitongsys/parquet-go/parquet"
	"github.com/xitongsys/parquet-go/writer"
)

// TestData is a struct for Parquet test data
type TestData struct {
	ID    int64  `parquet:"name=id, type=INT64"`
	Value int64  `parquet:"name=value, type=INT64"`
	Name  string `parquet:"name=name, type=BYTE_ARRAY, convertedtype=UTF8"`
}

func main() {
	// Create data directory if it doesn't exist
	dataDir := "./data"
	if err := os.MkdirAll(dataDir, 0755); err != nil {
		fmt.Printf("Failed to create data directory: %v\n", err)
		os.Exit(1)
	}

	// Generate CSV test data
	generateCSVData(filepath.Join(dataDir, "input.csv"))

	// Generate Parquet test data
	generateParquetData(filepath.Join(dataDir, "input.parquet"))

	fmt.Println("Test data generated successfully!")
}

func generateCSVData(filePath string) {
	// Create CSV file
	file, err := os.Create(filePath)
	if err != nil {
		fmt.Printf("Failed to create CSV file: %v\n", err)
		os.Exit(1)
	}
	defer file.Close()

	// Create CSV writer
	writer := csv.NewWriter(file)
	defer writer.Flush()

	// Write header
	if err := writer.Write([]string{"id", "value", "name"}); err != nil {
		fmt.Printf("Failed to write CSV header: %v\n", err)
		os.Exit(1)
	}

	// Write data
	data := [][]string{
		{"1", "100", "Alice"},
		{"2", "200", "Bob"},
		{"3", "300", "Charlie"},
		{"4", "400", "David"},
		{"5", "500", "Eve"},
	}

	for _, row := range data {
		if err := writer.Write(row); err != nil {
			fmt.Printf("Failed to write CSV row: %v\n", err)
			os.Exit(1)
		}
	}

	fmt.Printf("CSV data written to %s\n", filePath)
}

func generateParquetData(filePath string) {
	// Create Parquet file
	file, err := os.Create(filePath)
	if err != nil {
		fmt.Printf("Failed to create Parquet file: %v\n", err)
		os.Exit(1)
	}
	defer file.Close()

	// Create Parquet writer
	pw, err := writer.NewParquetWriter(file, new(TestData), 4)
	if err != nil {
		fmt.Printf("Failed to create Parquet writer: %v\n", err)
		os.Exit(1)
	}
	pw.CompressionType = parquet.CompressionCodec_SNAPPY

	// Write data
	data := []TestData{
		{ID: 1, Value: 100, Name: "Alice"},
		{ID: 2, Value: 200, Name: "Bob"},
		{ID: 3, Value: 300, Name: "Charlie"},
		{ID: 4, Value: 400, Name: "David"},
		{ID: 5, Value: 500, Name: "Eve"},
	}

	for _, row := range data {
		if err := pw.Write(row); err != nil {
			fmt.Printf("Failed to write Parquet row: %v\n", err)
			os.Exit(1)
		}
	}

	// Close writer
	if err := pw.WriteStop(); err != nil {
		fmt.Printf("Failed to close Parquet writer: %v\n", err)
		os.Exit(1)
	}

	fmt.Printf("Parquet data written to %s\n", filePath)
}
